#!/usr/bin/env python3
"""
EFinance数据源实现
基于efinance库的免费A股数据获取
"""

import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, Any, Optional
import time

from .base import BaseDataSource, DataSourceStatus
from ..utils.logger import get_logger

logger = get_logger("EFinanceSource")


class EFinanceSource(BaseDataSource):
    """EFinance数据源"""

    def __init__(self, config: Dict[str, Any] = None):
        super().__init__("efinance", config)
        self.ef = None
        self._initialize_fetcher()

    def _initialize_fetcher(self):
        """初始化数据获取器"""
        try:
            import efinance as ef
            self.ef = ef
            self.update_status(DataSourceStatus.AVAILABLE)
            logger.info("EFinance数据源初始化成功")
        except ImportError:
            self.update_status(DataSourceStatus.ERROR, "efinance未安装")
            logger.error("efinance未安装，请运行: pip install efinance")
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            logger.error(f"EFinance数据源初始化失败: {e}")

    def check_availability(self) -> bool:
        """检查数据源是否可用"""
        try:
            if not self.ef:
                self.update_status(DataSourceStatus.ERROR, "efinance未初始化")
                return False

            # 简单检查，不进行网络请求
            available = self.ef is not None
            self.update_status(
                DataSourceStatus.AVAILABLE if available else DataSourceStatus.ERROR,
                None if available else "efinance未可用"
            )
            return available
        except Exception as e:
            self.update_status(DataSourceStatus.ERROR, str(e))
            return False

    def get_stock_list(self) -> pd.DataFrame:
        """获取股票列表"""
        try:
            # 获取A股股票列表
            df = self.ef.stock.get_realtime_quotes()
            if not df.empty:
                # 转换为统一格式
                df = df[df['股票代码'].str.len() == 6]  # 只保留A股
                df = df.rename(columns={'股票代码': 'symbol', '股票名称': 'name'})
                df['symbol'] = df['symbol'].astype(str).str.zfill(6)
                df['data_source'] = 'efinance'
                # 添加市场信息
                df['market'] = df['symbol'].apply(lambda x: 'SH' if x.startswith('6') else 'SZ')
                return df[['symbol', 'name', 'market', 'data_source']]
        except Exception as e:
            logger.error(f"EFinance获取股票列表失败: {e}")
            return pd.DataFrame()

    def fetch_history_data(self, symbol: str, interval: str,
                          start_date: str, end_date: str) -> pd.DataFrame:
        """获取历史数据"""
        try:
            # 转换时间间隔
            freq_map = {
                '1d': '1',
                '1h': '60',
                '30m': '30',
                '15m': '15',
                '5m': '5'
            }
            freq = freq_map.get(interval, '1')

            # 获取历史数据
            df = self.ef.stock.get_quote_history(
                stock_codes=symbol,
                beg=start_date,
                end=end_date,
                klt=freq
            )

            if not df.empty:
                # 转换为统一格式
                df = df.rename(columns={
                    '日期': 'Date',
                    '开盘': 'Open',
                    '收盘': 'Close', 
                    '最高': 'High',
                    '最低': 'Low',
                    '成交量': 'Volume'
                })
                
                # 确保Date列是datetime类型
                df['Date'] = pd.to_datetime(df['Date'])
                
                # 选择需要的列并排序
                df = df[['Date', 'Open', 'High', 'Low', 'Close', 'Volume']].sort_values('Date')
                
                # 添加延迟避免频率限制
                time.sleep(0.1)
                
            return df

        except Exception as e:
            logger.error(f"EFinance获取历史数据失败 {symbol}: {e}")
            return pd.DataFrame()

    def get_stock_info(self, symbol: str) -> Dict[str, Any]:
        """获取股票信息"""
        try:
            # 获取股票基本信息
            df = self.ef.stock.get_base_info(symbol)
            if not df.empty:
                info = df.iloc[0].to_dict()
                return {
                    'symbol': symbol,
                    'name': info.get('股票名称', 'Unknown'),
                    'industry': info.get('所属行业', 'Unknown'),
                    'market': 'SH' if symbol.startswith('6') else 'SZ'
                }
            return {}
        except Exception as e:
            logger.error(f"EFinance获取股票信息失败 {symbol}: {e}")
            return {}
